#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2019 crane <crane@crane-pc>
#
# Distributed under terms of the MIT license.

"""
树的最长路径/diameter 应用: 在工程中求critical path(最短工期)
"""

from collections import defaultdict


class Diameter:
    def __init__(self, edges):
        self.edges = edges
        self.edge_map = defaultdict(dict)
        self.build_map()
        # self.diameter()

    def diameter(self):
        if not self.edges:
            return None, None, 0

        first_node = self.edges[0][0]

        # NOTE: bfs
        # end1, _ = self.bfs_longest_to(first_node)
        # end2, dis = self.bfs_longest_to(end1)

        # NOTE: dfs
        end1, _   = self.dfs_top(first_node)
        end2, dis = self.dfs_top(end1)

        return end1, end2, dis

    def build_map(self):
        for s,e,w in self.edges:
            self.edge_map[s][e] = w
            self.edge_map[e][s] = w

    def bfs_longest_to(self, node):
        queue = [(node, 0)]

        max_dis = 0
        max_node = None

        used = set()

        while queue:
            # print(queue)
            node, dis = queue.pop(0)
            used.add(node)

            for nei, weight in self.edge_map[node].items():
                if nei in used: continue

                # print(nei, weight)
                new_dis = dis + weight
                if new_dis > max_dis:
                    max_dis = new_dis
                    max_node = nei
                queue.append( (nei, new_dis) )

        return max_node, max_dis

    def dfs_top(self, node):
        self.used = set()
        return self.dfs_longest_to(node, 0)

    def dfs_longest_to(self, node, dis):
        # print(node, dis)
        self.used.add(node)
        max_end = node
        max_dis = dis

        for nei, weight in self.edge_map[node].items():
            if nei in self.used: continue
            new_end, new_dis = self.dfs_longest_to(nei, dis+weight)
            if new_dis > max_dis:
                max_dis = new_dis
                max_end = new_end

        # print('return ', max_end, max_dis)
        return max_end, max_dis


def build_case():
    edges = []
    edges.append( ('t', 'l', 5) )
    edges.append( ('t', 'r', 6) )

    edges.append( ('l', 'll', 100) )
    edges.append( ('l', 'lr', 200) )

    edges.append( ('r', 'rl', 40) )
    edges.append( ('r', 'rr', 100) )
    return edges


def test():
    edges = build_case()
    d = Diameter(edges)
    end1, end2, dia = d.diameter()
    print(end1, end2, dia)


def main():
    print("start main")
    test()

if __name__ == "__main__":
    main()
